package com.lucifer.cloud.boot.blog.config;

import jakarta.annotation.Resource;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingModel;
import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingOptions;
import org.springframework.ai.zhipuai.api.ZhiPuAiApi;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;
import java.util.Objects;

/**
 * @author lucifer
 * @date 2025/2/10 11:44
 */
@Configuration
@RequiredArgsConstructor
public class AiConfig {

    @Value("${spring.ai.model}")
    private String model;

    @Value("${spring.ai.zhipuai.api-key}")
    private String zhipuaiApiKey;

    @Value("${spring.ai.openai.api-key}")
    private String openaiApiKey;


    @Value("${spring.ai.openai.base-url}")
    private String openaiBaseUrl;

    @Value("${spring.datasource.pgvector.driverClassName}")
    private String driverClassName;

    @Value("${spring.datasource.pgvector.url}")
    private String url;

    @Value("${spring.datasource.pgvector.username}")
    private String username;

    @Value("${spring.datasource.pgvector.password}")
    private String password;

    @Bean(name = "embeddingModel")
    public EmbeddingModel embeddingModel() {
        ZhiPuAiApi zhiPuAiApi = new ZhiPuAiApi(zhipuaiApiKey);
        return new ZhiPuAiEmbeddingModel(zhiPuAiApi, MetadataMode.EMBED,
                ZhiPuAiEmbeddingOptions.builder()
                        .model("embedding-2")
                        .build());
    }

    @Bean(name = "openAiEmbeddingModel")
    public EmbeddingModel openAiEmbeddingModel() {

        OpenAiApi openAiApi = OpenAiApi.builder()
                .apiKey(openaiApiKey)
                .baseUrl(openaiBaseUrl)
                .build();

        OpenAiEmbeddingModel openAiEmbeddingModel = new OpenAiEmbeddingModel(
                openAiApi,
                MetadataMode.EMBED,
                OpenAiEmbeddingOptions.builder()
                        .model(OpenAiApi.EmbeddingModel.TEXT_EMBEDDING_ADA_002.value)
                        .user("user-6")
                        .build(),
                RetryUtils.DEFAULT_RETRY_TEMPLATE);
        return openAiEmbeddingModel;
    }




    @Resource
    private OllamaEmbeddingModel ollamaEmbeddingModel;

    @Bean(name = "pgJdbcTemplate")
    public JdbcTemplate jdbcTemplate() {
        DataSource dataSource = DataSourceBuilder.create().url(url)
                .username(username)
                .password(password)
                .driverClassName(driverClassName)
                .build();
        return new JdbcTemplate(dataSource);
    }



    @Bean(name = "pgVectorStore")
    public VectorStore vectorStore(@Qualifier("pgJdbcTemplate") JdbcTemplate pgJdbcTemplate) {
        PgVectorStore pgVectorStore = PgVectorStore.builder(pgJdbcTemplate, model())
                .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(PgVectorStore.PgIndexType.HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(false)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("vector_store")     // Optional: defaults to "vector_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
        return pgVectorStore;
    }



    @Bean(name = "chatMemoryVectorStore")
    public PgVectorStore chatMemoryVectorStore(@Qualifier("pgJdbcTemplate") JdbcTemplate pgJdbcTemplate) {
        PgVectorStore pgVectorStore = PgVectorStore.builder(pgJdbcTemplate, model())
                .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(PgVectorStore.PgIndexType.HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(false)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("chat_memory")     // Optional: defaults to "vector_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
        return pgVectorStore;
    }




    @Bean(name = "nl2sqlVectorStore")
    public PgVectorStore nl2sqlVectorStore(@Qualifier("pgJdbcTemplate") JdbcTemplate pgJdbcTemplate) {
        PgVectorStore pgVectorStore = PgVectorStore.builder(pgJdbcTemplate, model())
                .distanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(PgVectorStore.PgIndexType.HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(false)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("nl2sql_store")     // Optional: defaults to "nl2sql_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
        return pgVectorStore;
    }



    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }


    private EmbeddingModel model(){
        if(Objects.equals("zhipuai",model)){
            return embeddingModel();
        }else if (Objects.equals("ollama",model)){
            return ollamaEmbeddingModel;
        }else if (Objects.equals("openai",model)){
            return ollamaEmbeddingModel;
        }else {
            return ollamaEmbeddingModel;
        }
    }
}
